import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.naive_bayes import GaussianNB
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
import joblib


def train_and_evaluate_models(file_path):
    """
    直接从CSV文件读取数据,进行预处理，定义模型字典，训练并评估所有模型，
    并将结果保存到指定路径。
    """
    # 直接在函数内部指定数据文件路径

    # 读取数据
    # data = pd.read_csv(file_path)
    data = file_path
    healthcare_data = data.dropna()
    del healthcare_data['id']

    # 性别热编码
    gender_list = healthcare_data['gender']
    gender_hot = pd.get_dummies(gender_list)
    # 婚姻热编码
    ever_married_list = healthcare_data['ever_married']
    ever_married_hot = pd.get_dummies(ever_married_list)
    # 工作热编码
    work_type_list = healthcare_data['work_type']
    work_type_hot = pd.get_dummies(work_type_list)
    # 居住热编码
    residence_type_list = healthcare_data['Residence_type']
    residence_type_hot = pd.get_dummies(residence_type_list)
    # 吸烟热编码
    smoking_status_list = healthcare_data['smoking_status']
    smoking_status_hot = pd.get_dummies(smoking_status_list)

    del healthcare_data['gender'], healthcare_data['ever_married'], healthcare_data['work_type'], healthcare_data[
        'Residence_type'], healthcare_data['smoking_status']

    healthcare_data = pd.concat(
        [healthcare_data, gender_hot, ever_married_hot, work_type_hot, residence_type_hot, smoking_status_hot], axis=1)

    # X是特征矩阵，y是目标变量
    X, y = healthcare_data.drop('stroke', axis=1), healthcare_data['stroke']

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)
    # joblib.dump(X_train.columns, './health_model/health_columns.pkl')
    joblib.dump(X_train.columns, './models/health_model/health_columns.pkl')

    scaler = StandardScaler()
    print(X_train)
    X_train = scaler.fit_transform(X_train)
    X_test = scaler.transform(X_test)

    # joblib.dump(scaler, './health_model/health_scaler.pkl')
    joblib.dump(scaler, './models/health_model/health_scaler.pkl')
    print('health_scaler保存成功')


    # 定义模型字典
    models = {
        "logistic_regression": LogisticRegression(),
        "decision_tree": DecisionTreeClassifier(random_state=42),
        "random_forest": RandomForestClassifier(n_estimators=100, random_state=42),
        "Bayesian": GaussianNB(),
        "SVM": SVC(kernel='linear', probability=True, C=1),
        "KNN": KNeighborsClassifier(n_neighbors=5)
    }

    # 训练和评估模型
    for model_name, model in models.items():
        print(f"正在训练和评估 {model_name}...")
        model.fit(X_train, y_train)
        y_pred = model.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        print(f"{model_name} 准确率: {accuracy}")
        print("分类报告:")
        print(classification_report(y_test, y_pred))
        print("混淆矩阵:")
        print(confusion_matrix(y_test, y_pred))
        # 保存模型
        model_file_path = f'./models/health_model/{model_name}.pkl'
        joblib.dump(model, model_file_path)
        print(f"{model_name}模型已保存")


# 主程序调用
if __name__ == "__main__":
    file_path = "../data/healthcare-dataset-stroke-data.csv"
    train_and_evaluate_models(file_path)
